import gradio as gr
import requests
from utils.utils import gen_random_job_id, logger, base64_to_audio_file, audio_to_base64
import os

# 创建临时目录
if not os.path.exists("tmp"):
    os.makedirs("tmp")

api_url = "http://localhost:6006"

class AudioDH:
    def __init__(self):
        self.system_prompt = "你是minglog，是由骆明开发的聊天机器人。你擅长编写人工智能和大模型相关的代码，有丰富的项目开发经验。"

    def asr(self, audio_path):
        api = f"{api_url}/asr"
        logger.info(f"ASR In: {audio_path}")
        resp = requests.post(api, json={"audio_path": audio_to_base64(audio_path)})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            text = resp.json().get("content")
            logger.info(f"ASR Out: {text}")
            print("-" * 20)
        else:
            raise gr.Error("ASR 出错，请检查ASR服务。")
        return text

    def llm(self, text, system_message, message):
        api = f"{api_url}/llm"
        logger.info(f"LLM In: {text}")
        resp = requests.post(api, json={"text": text, "system_message": system_message, "message": message})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            resp_text = resp.json().get("content")
            logger.info(f"LLM Out: {resp_text}")
            print("-" * 20)
        else:
            raise gr.Error("LLM 出错，请检查LLM服务。")
        return resp_text

    def tts(self, text):
        api = f"{api_url}/tts"
        logger.info(f"TTS In: {text}")
        resp = requests.post(api, json={"text": text})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            audio_base64 = resp.json().get("content")
            save_path = f"tmp/{gen_random_job_id()}.mp3"
            base64_to_audio_file(audio_base64, save_path)
            logger.info(f"TTS Out: {save_path}")
            print("-" * 20)
            return save_path
        else:
            raise gr.Error("TTS 出错，请检查TTS服务。")
    
    def respond(self, audio, history, chat_ui):
        text = self.asr(audio)

        # 初始化对话上下文
        if history is None:
            history = []

        if chat_ui is None:
            chat_ui = []
        llm_text = self.llm(text, self.system_prompt, history)[0]
        # 更新对话历史
        history.append({"role": "user", "content": text})
        history.append({"role": "assistant", "content": llm_text})

        return_audio = self.tts(llm_text)
        user_html = f"<p>🧑‍💬： {text}</p><audio controls src='http://localhost:7860/gradio_api/file={audio}'></audio>"
        bot_html = f"<p>🤖： {llm_text}</p><audio controls autoplay src='http://localhost:7860/gradio_api/file={return_audio}'></audio>"

        # 更新 UI 显示历史
        chat_ui.append((user_html, bot_html))
        return chat_ui, history, chat_ui, None
    
    def create_ui(self):
        with gr.Blocks() as demo:
            chatbot = gr.Chatbot()
            audio_input = gr.Audio(sources="microphone", type="filepath", label="上传音频")
            history_state = gr.State(value=[])  # 存储每个用户的对话历史
            chat_ui_state = gr.State(value=[])  # 页面显示的内容
            audio_input.stop_recording(fn=self.respond, inputs=[audio_input, history_state, chat_ui_state], outputs=[chatbot, history_state, chat_ui_state, audio_input])
            # audio_input.upload(fn=self.respond, inputs=[audio_input, history_state, chat_ui_state], outputs=[chatbot, history_state, chat_ui_state, audio_input])
        return demo

audio_dh = AudioDH()
demo = audio_dh.create_ui()
demo.launch(allowed_paths=["tmp"], server_name="0.0.0.0", server_port=7860)
